#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MQTT服务程序
功能：
1. 监听指定的多个MQTT topic
2. 等待所有topic的消息都到达后（通过时间戳判断）再一起处理
3. 延时后推送到指定的topic
4. 提供Flask健康检测接口
"""

import os
import json
import time
import threading
import uuid
from typing import List, Dict, Any, Optional
import logging
from dotenv import load_dotenv
import paho.mqtt.client as mqtt
from flask import Flask, jsonify
from collections import defaultdict, deque
from predict import Model_Simulation

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class MessageBatch:
    """消息批次类，用于管理同一时间戳的多路消息"""
    
    def __init__(self, timestamp: float, batch_id: str):
        self.timestamp = timestamp
        self.batch_id = batch_id
        self.messages = {}  # topic -> message_data
        self.created_time = time.time()
        self.is_complete = False
    
    def add_message(self, topic: str, message_data: dict):
        """添加消息到批次"""
        self.messages[topic] = message_data
        logger.info(f"批次 {self.batch_id} 收到来自 {topic} 的消息")
    
    def is_ready(self, required_topics: List[str]) -> bool:
        """检查批次是否完整（所有必需的topic都有消息）"""
        return all(topic in self.messages for topic in required_topics)
    
    def get_all_messages(self) -> Dict[str, dict]:
        """获取批次中的所有消息"""
        return self.messages.copy()
    
    def get_age(self) -> float:
        """获取批次年龄（秒）"""
        return time.time() - self.created_time

class MQTTService:
    def __init__(self):
        # MQTT配置
        self.mqtt_host = os.getenv('MQTT_HOST', 'localhost')
        self.mqtt_port = int(os.getenv('MQTT_PORT', 1883))
        self.mqtt_username = os.getenv('MQTT_USERNAME', '')
        self.mqtt_password = os.getenv('MQTT_PASSWORD', '')
        self.mqtt_client_id = os.getenv('MQTT_CLIENT_ID', f'mqtt_test_service_{uuid.uuid4().hex[:8]}')
        
        # Topic配置
        self.listen_topics = os.getenv('LISTEN_TOPICS', 'test/input').split(',')
        self.publish_topic = os.getenv('PUBLISH_TOPIC', 'test/output')
        
        # 延时配置（秒）- 已取消延时
        self.delay_seconds = 0.0
        # 批次超时配置（秒）
        self.batch_timeout = float(os.getenv('BATCH_TIMEOUT', '10.0'))
        
        # 时间戳容差（秒）- 用于判断消息是否属于同一批次
        self.timestamp_tolerance = float(os.getenv('TIMESTAMP_TOLERANCE', '2.0'))
        
        # Flask配置
        self.flask_host = os.getenv('FLASK_HOST', '0.0.0.0')
        self.flask_port = int(os.getenv('FLASK_PORT', 5000))
        
        # 状态统计
        self.stats = {
            'messages_received': 0,
            'messages_published': 0,
            'batches_created': 0,
            'batches_completed': 0,
            'batches_timeout': 0,
            'last_message_time': None,
            'service_start_time': time.time(),
            'is_connected': False
        }
        
        # 初始化MQTT客户端
        self.client = mqtt.Client(client_id=self.mqtt_client_id)
        self.client.on_connect = self.on_connect
        self.client.on_message = self.on_message
        self.client.on_disconnect = self.on_disconnect
        
        # 设置认证信息
        if self.mqtt_username and self.mqtt_password:
            self.client.username_pw_set(self.mqtt_username, self.mqtt_password)
        
        # 初始化Flask应用
        self.app = Flask(__name__)
        self.setup_flask_routes()
        
        # 消息批次管理
        self.message_batches = {}  # batch_id -> MessageBatch
        self.batch_lock = threading.Lock()
        
        # 消息队列（用于延时处理）
        self.message_queue = []
        self.queue_lock = threading.Lock()

        # 预测模型初始化（启动时初始化一次）
        self.model = Model_Simulation({})
        # 首次带有config的批次将触发一次性按config重建模型
        self.model_initialized_from_config = False

    
    def setup_flask_routes(self):
        """设置Flask路由"""
        @self.app.route('/health', methods=['GET'])
        def health_check():
            """健康检测接口"""
            return jsonify({
                'status': 'healthy',
                'timestamp': time.time(),
                'stats': self.stats,
                'mqtt_connected': self.stats['is_connected'],
                'uptime_seconds': time.time() - self.stats['service_start_time']
            })
        
        @self.app.route('/stats', methods=['GET'])
        def get_stats():
            """获取统计信息"""
            return jsonify(self.stats)
        
        @self.app.route('/topics', methods=['GET'])
        def get_topics():
            """获取topic配置信息"""
            return jsonify({
                'listen_topics': self.listen_topics,
                'publish_topic': self.publish_topic,
                'batch_timeout': self.batch_timeout,
                'timestamp_tolerance': self.timestamp_tolerance
            })
        
        @self.app.route('/batches', methods=['GET'])
        def get_batches():
            """获取当前批次信息"""
            with self.batch_lock:
                batch_info = {}
                for batch_id, batch in self.message_batches.items():
                    batch_info[batch_id] = {
                        'timestamp': batch.timestamp,
                        'age_seconds': batch.get_age(),
                        'messages_count': len(batch.messages),
                        'is_complete': batch.is_ready(self.listen_topics),
                        'topics': list(batch.messages.keys())
                    }
                return jsonify(batch_info)
    
    def get_next_batch_id(self) -> str:
        """获取下一个批次ID"""
        return f"batch_{uuid.uuid4().hex[:8]}_{int(time.time())}"
    
    def find_or_create_batch(self, message_timestamp: float) -> tuple[str, MessageBatch]:
        """查找或创建合适的消息批次"""
        with self.batch_lock:
            # 查找时间戳相近的现有批次
            for batch_id, batch in self.message_batches.items():
                if abs(batch.timestamp - message_timestamp) <= self.timestamp_tolerance:
                    return batch_id, batch
            
            # 创建新批次
            batch_id = self.get_next_batch_id()
            new_batch = MessageBatch(message_timestamp, batch_id)
            self.message_batches[batch_id] = new_batch
            self.stats['batches_created'] += 1
            logger.info(f"创建新批次 {batch_id}，时间戳: {message_timestamp}")
            
            return batch_id, new_batch
    
    def cleanup_expired_batches(self):
        """清理超时的批次"""
        current_time = time.time()
        expired_batches = []
        
        with self.batch_lock:
            for batch_id, batch in self.message_batches.items():
                if batch.get_age() > self.batch_timeout:
                    expired_batches.append(batch_id)
        
        # 处理超时批次
        for batch_id in expired_batches:
            with self.batch_lock:
                if batch_id in self.message_batches:
                    batch = self.message_batches.pop(batch_id)
                    self.stats['batches_timeout'] += 1
                    logger.warning(f"批次 {batch_id} 超时，丢弃 {len(batch.messages)} 条消息")
    
    def on_connect(self, client, userdata, flags, rc):
        """MQTT连接回调"""
        if rc == 0:
            logger.info(f"成功连接到MQTT服务器 {self.mqtt_host}:{self.mqtt_port}")
            self.stats['is_connected'] = True
            
            # 订阅监听topic
            for topic in self.listen_topics:
                topic = topic.strip()
                if topic:
                    client.subscribe(topic)
                    logger.info(f"已订阅topic: {topic}")
        else:
            logger.error(f"MQTT连接失败，返回码: {rc}")
            self.stats['is_connected'] = False
    
    def on_disconnect(self, client, userdata, rc):
        """MQTT断开连接回调"""
        logger.warning(f"MQTT连接断开，返回码: {rc}")
        self.stats['is_connected'] = False
    
    def on_message(self, client, userdata, msg):
        """MQTT消息接收回调"""
        try:
            topic = msg.topic
            payload = msg.payload.decode('utf-8')
            
            logger.info(f"收到消息 - Topic: {topic}, Payload: {payload}")
            self.stats['messages_received'] += 1
            self.stats['last_message_time'] = time.time()
            
            # 解析消息，提取时间戳
            try:
                message_data = json.loads(payload)
                message_timestamp = message_data.get('timestamp', time.time())
                # 兼容毫秒时间戳
                if isinstance(message_timestamp, (int, float)) and message_timestamp > 1e11:
                    message_timestamp = message_timestamp / 1000.0
            except (json.JSONDecodeError, KeyError):
                # 如果消息格式不正确，使用当前时间
                message_timestamp = time.time()
                message_data = {'raw_payload': payload, 'timestamp': message_timestamp}
            
            # 查找或创建批次
            batch_id, batch = self.find_or_create_batch(message_timestamp)
            
            # 添加消息到批次
            batch.add_message(topic, message_data)
            
            # 检查批次是否完整
            if batch.is_ready(self.listen_topics):
                logger.info(f"批次 {batch_id} 完整，准备处理")
                self.stats['batches_completed'] += 1

                # 直接发布（取消延时和队列）
                batch_data = {
                    'batch_id': batch_id,
                    'timestamp': batch.timestamp,
                    'messages': batch.get_all_messages(),
                }
                self.publish_batch(batch_data)

                # 从活跃批次中移除
                with self.batch_lock:
                    if batch_id in self.message_batches:
                        del self.message_batches[batch_id]

            # 每次消息后做一次过期清理
            self.cleanup_expired_batches()
            
        except Exception as e:
            logger.error(f"处理消息时出错: {e}")
    
    # 取消队列与后台处理线程相关方法
    def queue_batch(self, batch_id: str, batch: MessageBatch):
        logger.debug("queue_batch 已弃用（延时取消）")
        pass
    
    def process_message_queue(self):
        logger.debug("process_message_queue 已弃用（延时取消）")
        pass
    
    def publish_batch(self, batch_data: dict):
        """发布完整批次到指定topic：合并pointData，调用predict并新增output字段"""
        try:
            # 1) 合并pointData与config
            merged_point, merged_config, base_msg = self._merge_batch_messages(batch_data['messages'])

            # 1.1) 如首次收到config，则按config重建模型（一次性）
            if (not self.model_initialized_from_config) and isinstance(merged_config, dict) and len(merged_config) > 0:
                try:
                    self.model = Model_Simulation(merged_config)
                    self.model_initialized_from_config = True
                    logger.info("首次按config初始化模型完成，后续不再重复初始化")
                except Exception as e:
                    logger.error(f"按config初始化模型失败: {e}")

            # 2) 构建模型输入（直接使用pointData字典）
            model_input = merged_point if isinstance(merged_point, dict) else {}

            # 3) 调用模型
            output = {}
            try:
                logger.info(f"开始模型仿真")
                # output = self.model.process_step(model_input)
                output = self.model.mockData(model_input)
                logger.info(f"模型仿真完成")
            except Exception as e:
                logger.error(f"模型计算失败: {e}")
                output = {'error': f'model_failed: {str(e)}'}

            # 不再根据config覆盖学习率或其他参数（初始化阶段已处理）

            # 4) 生成发布消息：保持结构一致，新增output
            publish_msg = {}
            if isinstance(base_msg, dict):
                publish_msg.update(base_msg)
            # 仅当合并结果非空时覆盖；否则继承原始 pointData
            if isinstance(merged_point, dict) and len(merged_point) > 0:
                publish_msg['pointData'] = merged_point
            else:
                if isinstance(base_msg, dict) and isinstance(base_msg.get('pointData'), dict):
                    publish_msg['pointData'] = base_msg['pointData']
                else:
                    publish_msg['pointData'] = {}
            if merged_config is not None:
                publish_msg['config'] = merged_config
            publish_msg['output'] = self._to_jsonable(output)
            message = json.dumps(publish_msg, ensure_ascii=False)
            # 使用UTF-8字节发布，避免编码歧义
            result = self.client.publish(self.publish_topic, message.encode('utf-8'))

            if result.rc == mqtt.MQTT_ERR_SUCCESS:
                logger.info(f"批次 {batch_data['batch_id']} 已发布到 {self.publish_topic}: 聚合消息")
                self.stats['messages_published'] += 1
            else:
                logger.error(f"发布批次失败，返回码: {result.rc}")

        except Exception as e:
            logger.error(f"发布批次时出错: {e}")

    def _merge_batch_messages(self, messages_by_topic: dict):
        """合并批次内多路消息的pointData与config，并返回一个基准消息副本"""
        merged_point = {}
        merged_config = {}
        base_msg = None
        for _, msg in messages_by_topic.items():
            if isinstance(msg, dict) and base_msg is None:
                base_msg = msg.copy()
            # 合并pointData
            pd = msg.get('pointData') if isinstance(msg, dict) else None
            if isinstance(pd, dict):
                merged_point.update(pd)
            # 合并config
            cfg = msg.get('config') if isinstance(msg, dict) else None
            # 兼容字符串形式的JSON配置
            if isinstance(cfg, str):
                try:
                    cfg = json.loads(cfg)
                except Exception:
                    pass
            if isinstance(cfg, dict):
                merged_config.update(cfg)
        return merged_point, (merged_config if merged_config else None), (base_msg or {})

    def _build_model_input(self, merged_point: dict) -> dict:
        """直接返回合并后的pointData（模型内部自行取用）"""
        return merged_point if isinstance(merged_point, dict) else {}

    def _to_jsonable(self, obj: Any) -> Any:
        """将可能包含numpy类型的对象转换为可JSON序列化的原生类型"""
        try:
            import numpy as _np  # 局部导入，避免硬依赖
        except Exception:
            _np = None

        if obj is None:
            return None
        if isinstance(obj, (str, int, float, bool)):
            return obj
        if _np is not None and isinstance(obj, (_np.generic,)):
            return obj.item()
        if isinstance(obj, dict):
            return {self._to_jsonable(k): self._to_jsonable(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple)):
            return [self._to_jsonable(v) for v in obj]
        try:
            return float(obj)  # 最后尝试转为float
        except Exception:
            return str(obj)
    
    def start(self):
        """启动服务"""
        try:
            # 取消队列与清理后台线程，改为同步处理
            logger.info("延时与后台队列已取消，采用同步处理批次")
            
            # 连接MQTT服务器
            logger.info(f"正在连接MQTT服务器 {self.mqtt_host}:{self.mqtt_port}")
            self.client.connect(self.mqtt_host, self.mqtt_port, 60)
            
            # 启动MQTT客户端循环
            mqtt_thread = threading.Thread(target=self.client.loop_forever, daemon=True)
            mqtt_thread.start()
            logger.info("MQTT客户端循环已启动")
            
            # 启动Flask应用
            logger.info(f"正在启动Flask健康检测服务 {self.flask_host}:{self.flask_port}")
            self.app.run(
                host=self.flask_host,
                port=self.flask_port,
                debug=False,
                use_reloader=False
            )
            
        except KeyboardInterrupt:
            logger.info("收到中断信号，正在关闭服务...")
            self.stop()
        except Exception as e:
            logger.error(f"启动服务时出错: {e}")
            self.stop()
    
    def stop(self):
        """停止服务"""
        try:
            if self.client.is_connected():
                self.client.disconnect()
            logger.info("MQTT服务已停止")
        except Exception as e:
            logger.error(f"停止服务时出错: {e}")

def main():
    """主函数"""
    logger.info("正在启动MQTT多路监听测试服务...")
    
    # 打印配置信息
    service = MQTTService()
    logger.info(f"MQTT配置: {service.mqtt_host}:{service.mqtt_port}")
    logger.info(f"监听Topic: {service.listen_topics}")
    logger.info(f"发布Topic: {service.publish_topic}")
    logger.info(f"批次超时: {service.batch_timeout} 秒")
    logger.info(f"时间戳容差: {service.timestamp_tolerance} 秒")
    logger.info(f"健康检测端口: {service.flask_host}:{service.flask_port}")
    
    # 启动服务
    service.start()

if __name__ == "__main__":
    main()
